# 深度学习
import torch
import torch.nn as nn
import torch.nn.functional as F

class DigitClassifier(nn.Module):
    def __init__(self):
        super(DigitClassifier, self).__init__()
        
        # 第一层：捕获低级特征（边缘、纹理等），使用较多通道
        self.conv1 = nn.Conv2d(
            in_channels=1,        # 输入通道（MNIST是灰度图，通道数为1）
            out_channels=32,      # 大量通道捕获低级特征
            kernel_size=3,
            padding=1             # 保持尺寸不变
        )
        
        # 第二层：组合低级特征为中层特征，减少通道数
        self.conv2 = nn.Conv2d(
            in_channels=32,       # 继承上一层的32个通道
            out_channels=16,      # 减少通道，组合特征
            kernel_size=3,
            padding=1
        )
        
        # 第三层：提取高级抽象特征，进一步减少通道
        self.conv3 = nn.Conv2d(
            in_channels=16,       # 继承上一层的16个通道
            out_channels=8,       # 更少的通道，专注高级特征
            kernel_size=3,
            padding=1
        )
        
        # 全连接层：映射到10个分类（0-9数字）
        self.fc = nn.Linear(8 * 28 * 28, 10)  # 输入维度为8个通道的特征图展平

    def forward(self, x):
        # 输入形状：[B, 1, 28, 28]
        
        # 第一层卷积 + ReLU激活
        x = F.relu(self.conv1(x))  # 输出形状：[B, 32, 28, 28]
        
        # 第二层卷积 + ReLU激活
        x = F.relu(self.conv2(x))  # 输出形状：[B, 16, 28, 28]
        
        # 第三层卷积 + ReLU激活
        x = F.relu(self.conv3(x))  # 输出形状：[B, 8, 28, 28]
        
        # 展平特征图
        x = x.view(-1, 8 * 28 * 28)  # 形状：[B, 8*28*28]
        
        # 全连接层，输出10个类别得分
        x = self.fc(x)  # 形状：[B, 10]
        
        return x